package book;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.PriorityQueue;
import java.util.Set;

// 边权类，同时实现比较接口，用于排序
class Edge implements Comparable<Edge> {
    public int from;
    public int to;
    public int weight;

    Edge(int f, int t, int w) {
        from = f;
        to = t;
        weight = w;
    }

    @Override
    public int compareTo(Edge o) {
        // 排序依据：根据边权从小到大
        return Integer.compare(weight, o.weight);
    }
}

// 工具类 - 并查集
class UnionFind {
    int parent[], n;

    UnionFind(int n) {
        this.n = n;
        parent = new int[n];
        for (int i = 0; i < n; i++) {
            parent[i] = i;
        }
    }

    void pair(int index1, int index2) {
        parent[find(index2)] = find(index1);
    }

    int find(int index) {
        if (parent[index] != index) {
            parent[index] = find(parent[index]);
        }
        return parent[index];
    }

    boolean check(int index1, int index2) {
        return find(index1) == find(index2);
    }
}

public class MSTSolution {
    // Prim算法求最小生成树
    static void prim(int n, int[][] edges) {
        // 在Prim算法中使用邻接表存储图信息
        List<List<Edge>> wg = new ArrayList<>(n);
        for (int i = 0; i < n; i++)
            wg.add(new ArrayList<>());
        for (int[] es : edges) {
            // 无向图转邻接表时两个方向都要添加
            Edge e1 = new Edge(es[0], es[1], es[2]);
            wg.get(e1.from).add(e1);
            Edge e2 = new Edge(es[1], es[0], es[2]);
            wg.get(e2.from).add(e2);
        }

        // 使用集合维护当前已加入最小生成树的顶点
        Set<Integer> vts = new HashSet<>();

        // 使用优先队列快速获取当前最短的边权
        PriorityQueue<Edge> pq = new PriorityQueue<>(edges.length);

        // 最终生成的最小生成树用边权集表示
        List<Edge> mst = new ArrayList<>(n);

        // 初始化：从顶点0开始构建最小生成树
        vts.add(0);
        for (Edge e : wg.get(0))
            pq.offer(e);

        // 循环执行Prim算法：不断将当前符合条件的最短边加入MST
        while (vts.size() < n && !pq.isEmpty()) {
            Edge minEdge = pq.poll();
            // 若当前边的终点已经在MST中，该边无效，跳过
            if (vts.contains(minEdge.to)) 
                continue;

            // 否则，将这条边加入到MST中
            // 因为当前边是从优先队列中取出，所以保证是所有邻接边中最短的一条
            vts.add(minEdge.to);
            mst.add(minEdge);

            // 往外扩展：将新顶点能访问到的边加入优先队列
            for (Edge e : wg.get(minEdge.to))
                if (!vts.contains(e.to))
                    pq.offer(e);
        }

        // 完成，输出结果
        System.out.println("Prim:");
        if (vts.size() == n) {
            for (Edge e : mst)
                System.out.println(e.from + "->" + e.to);
        } else {
            System.out.println("当前图不连通！");
        }
    }

    // Kruskal算法求最小生成树
    static void kruskal(int n, int[][] edges) {
        // 直接用边权集表示图，并按边权排序
        Edge[] sortedEdges = new Edge[edges.length];
        for (int i = 0; i < edges.length; i++)
            sortedEdges[i] = new Edge(edges[i][0], edges[i][1], edges[i][2]);
        // Edge类实现了Comparable接口后，可以直接排序
        Arrays.sort(sortedEdges);
            
        // 使用并查集快速判断两个顶点是否已经在同一棵子树当中
        UnionFind unionFind = new UnionFind(n);

        // 最终生成的最小生成树用边权集表示
        List<Edge> mst = new ArrayList<>(n);

        // 循环执行Kruskal算法，不断将符合条件的最短边加入MST当中
        for (Edge minEdge : sortedEdges) {
            // 如果最小生成树构建成功则退出循环
            if (mst.size() == n - 1)
                break;

            // 判断该条边两个顶点是否连通
            // 如果已经连通，不可再加入MST(否则会形成环)
            if (!unionFind.check(minEdge.from, minEdge.to)) {
                // 将新的边加入MST，同时在并查集标记两个顶点连通
                unionFind.pair(minEdge.from, minEdge.to);
                mst.add(minEdge);
            }
        }

        // 完成，输出结果
        System.out.println("Kruskal:");
        if (mst.size() == n - 1) {
            for (Edge e : mst)
                System.out.println(e.from + "->" + e.to);
        } else {
            System.out.println("当前图不连通！");
        }
    }

    public static void main(String[] args) {
        int n = 6;
        int[][] allEdges = {{0, 1, 34}, {0, 5, 19}, {0, 2, 46}, {2, 5, 25}, {2, 3, 17}, {3, 4, 38}, {3, 5, 25}, {4, 5, 26}, {4, 1, 12}};
        prim(n, allEdges);
        kruskal(n, allEdges);
    }
}


